Код программ отличается от естественного языка из-за его формализма и строгости, однако ничто не мешает воспринимать его как последовательность токенов и работать с ним, как с обычным языком. Существуют исследования, которые показали, что модель BERT, обученная на большом наборе данных, неплохо справляется с некоторыми задачами, связанными с обработкой программного кода. В этом посте я буду решать задачу автогенерации комментариев к нему. Вы узнаете, как подготовить данные для обучения, настроить нейросеть и получить результат.

Данные

Данные представлены в виде набора пар [функция — комментарий] для различных языков программирования (awesome Code Search Net Challenge dataset). Кстати говоря, этот набор изначально был создан не для этой задачи, однако его можно легко перепрофилировать под свои нужды. 

Я не стану очищать данные, это описано здесь. Я же буду использовать уже предварительно обработанные данные в объеме 1 % от общего количества образцов в наборе, так как обучение модели занимает довольно много времени. Но, как можно будет убедиться в будущем, генерация комментариев даже на 1 % данных выглядит неплохо. Если у вас есть время и ресурсы, можете обучить модель на всём наборе и получить результаты получше.

CodeBERT

Предварительно обученная модель, которую я буду использовать, взята из статьи исследовательского подразделения Microsoft. В этой модели также использовался набор данных CodeSearchNet, но вместо генерирования комментариев он использовался для обучения модели на основе RoBERTa удобному для восприятия представлению кода и естественного языка. Использование больших языковых моделей для представления текста удобным способом в настоящее время является обычной практикой, поскольку они показали свою эффективность для решения других задач. 

Загрузка, установка и импортирование библиотек

!pip install transformers !git clone -q https://github.com/microsoft/CodeXGLUE.git   import json from dataclasses import dataclass import numpy as np import pandas as pd from transformers import AutoTokenizer

Здесь я прописываю пути до файлов с данными и оборачиваю их в структуру для более удобного дальнейшего использования:

PATH_TO_TRAIN_DATA = '/content/train.csv' PATH_TO_TEST_DATA = '/content/test.csv' PATH_TO_VALIDATION_DATA = '/content/validation.csv' #validation, test and train data_struct = { 'train' : pd.read_csv(PATH_TO_TRAIN_DATA), 'test' : pd.read_csv(PATH_TO_TEST_DATA), 'valid' : pd.read_csv(PATH_TO_VALIDATION_DATA) }

Инициализирую две вспомогательные функции: токенизации текста и записи DataFrame в JSON-файл, так как именно в таком формате требуется подавать данные для модели.

} def write_into_json_file(json_file_name: str, data: pd.DataFrame): ''' json_file_name - name output json file data - pandas data frame write your pandas data to json file ''' with open(json_file_name, 'w') as current_file: for index, current_row in data.iterrows(): current_file.write(json.dumps(current_row.to_dict()) + '\n')   def split_data(split_column: str, new_column: str, data: pd.DataFrame)-> pd.DataFrame: ''' split items in column data - your pandas data frame split_column - column in your pd.df ''' data[new_column] = data[split_column].apply(lambda current_item: current_item.split()) return data

Реализую небольшую предобработку данных с помощью функций, описанных выше:

#preproc data for type_data, value in data_struct.items(): #split target colums code_tokens_step = split_data('code', 'code_tokens', value) docs_tokens_step = split_data('comment', 'docstring_tokens', code_tokens_step) data_struct[type_data] = docs_tokens_step #create json file write_into_json_file(f'/content/{type_data}.jsonl', data_struct[type_data])

Далее создаю конфигурационный класс для модели и на его основе прописываю всю конфигурацию:

@dataclass class ConfigurationModel: learning_rate : float  batch_size : int  beam : int test_file : str source_size : int target_size : int path_to_data_directory : str path_to_output_data_directory : str train_file : str dev_file : str count_epochs : int pretrained_model : str   configuration_codetext_model = ConfigurationModel( learning_rate = 5e-5, batch_size = 8, beam = 10, source_size = 256, target_size = 512, path_to_data_directory = '.', path_to_output_data_directory = 'model_for_java', train_file = '/content/train.jsonl', dev_file = '/content/valid.jsonl', test_file = '/content/test.jsonl', count_epochs = 10, pretrained_model = 'microsoft/codebert-base', ) configuration_codetext_model 

Обучение

Теперь, когда данные обработаны и представлены в удобном формате, можно приступать к обучению. Сделаю этоОбучу модель на обучающей выборке. В качестве метрики использую BLEU-4 (четвёрка означает, что количество словесных n-gram = 4), которая распределена от 0 до 1, но в нашем примере будет использоваться BLEU-4 * 100%. Эта метрика используется в задачах машинного перевода, но и для генерации текста она также хорошо подходит. Если брать задачи машинного перевода, то даже для человека bleu = [0.6:0.7] — отличный результат, потому что каждый человек может перевести текст по-разному. Точности в единицу достигнуть почти невозможно. 

Если посмотреть на исходную задачу, то, во-первых, модель должна сгенерировать текст, а во-вторых, это не просто текст, а осмысленный комментарий к коду. Поэтому ожидать больших значений метрики bleu не стоит. 

#run train model !python /content/CodeXGLUE/Code-Text/code-to-text/code/run.py \ --do_train \ --do_eval \ --do_lower_case \ --model_type roberta \ --model_name_or_path {configuration_codetext_model.pretrained_model} \ --train_filename {configuration_codetext_model.train_file} \ --dev_filename {configuration_codetext_model.dev_file} \ --output_dir {configuration_codetext_model.path_to_output_data_directory} \ --max_source_length {configuration_codetext_model.source_size} \ --max_target_length {configuration_codetext_model.target_size} \ --beam_size {configuration_codetext_model.beam} \ --train_batch_size {configuration_codetext_model.batch_size} \ --eval_batch_size {configuration_codetext_model.batch_size} \ --learning_rate {configuration_codetext_model.learning_rate} \ --num_train_epochs {configuration_codetext_model.count_epochs}

Обучение

После обучения модели её можно проверить на отдельной выборке:

binary_model_file = '/content/model_for_java/checkpoint-best-bleu/pytorch_model.bin' !python  /content/CodeXGLUE/Code-Text/code-to-text/code/run.py \     --do_test \     --model_type roberta \     --model_name_or_path microsoft/codebert-base \     --load_model_path {binary_model_file} \     --dev_filename {configuration_codetext_model.dev_file} \     --test_filename {configuration_codetext_model.test_file} \     --output_dir {configuration_codetext_model.path_to_output_data_directory} \     --max_source_length {configuration_codetext_model.source_size} \     --max_target_length {configuration_codetext_model.target_size} \     --beam_size {configuration_codetext_model.beam} \     --eval_batch_size {configuration_codetext_model.batch_size}

Как можно увидеть, bleu-4 = 11, и это неплохо для такой задачи, даже с учётом того, что bleu в нашем случае распределена от 0 до 100.

Далее считаю получившиеся результаты:

path_to_gold = '/content/model_for_java/test_1.gold' path_to_output = '/content/model_for_java/test_1.output'

Инициализирую функцию считывания из txt-файла:

def read_result_txt_file(txt_file: str)-> list: with open(txt_file) as file:return [' '.join(line.rstrip().replace('\t', ' ').split(' ')[1:]) for line in file]

И для удобства считаю всё в DataFrame:

def read_result_txt_file(txt_file: str)-> list: #true comments and predicted true_sent = read_result_txt_file(path_to_gold) pred_sent = read_result_txt_file(path_to_output) result_data_frame = pd.DataFrame( { 'code' : data_struct['test']['code'], 'true' : true_sent, 'pred' : pred_sent } )   result_data_frame.head(10)

Вывод 10 примеров кода, оригинальных комментариев и комментариев, сгенерированных моделью.

Теперь попробую субъективно сравнить оригинальный комментарий со сгенерированным по шкале от 1 до 5. Code — исходный код, true — исходный комментарий, pred — сгенерированный.

Пример 1:

Code: public t includeas(final class template) { blacklist = false; string[] properties = getallbeanpropertynames(template, false); include(properties); return _this(); }

True: defines included property names as public properties of given template class. sets to black list mode.

Pred: create a new resource

Оценка: 1 — абсолютно непонятно, о чём идёт речь.

Пример 2:

Code: int setdirect(int v0, int v1, int v2, int v3, int v4) { return offset + v0*stride0 + v1*stride1 + v2*stride2 + v3*stride3 + v4*stride4; }

True: experimental : should be package private

Pred: sets the value for the specified point.

Оценка: 4 — исходный комментарий абсолютно никак не отражает функциональность, в отличие от сгенерированного. 

Пример 3:

Code: static private servicetype checkifdap4(string location) throws ioexception { // strip off any trailing dap4 prefix if (location.endswith(«.dap»)) location = location.substring(0, location.length() — «.dap».length()); else if (location.endswith(«.dmr»)) location = location.substring(0, location.length() — «.dmr».length()); else if (location.endswith(«.dmr.xml»)) location = location.substring(0, location.length() — «.dmr.xml».length()); else if (location.endswith(«.dsr»)) location = location.substring(0, location.length() — «.dsr».length()); try (httpmethod method = httpfactory.get(location + «.dmr.xml»)) { int status = method.execute(); if (status == 200) { header h = method.getresponseheader(«content-type»); if ((h != null) && (h.getvalue() != null)) { string v = h.getvalue(); if (v.startswith(«application/vnd.opendap.org»)) return servicetype.dap4; } } if (status == httpstatus.sc_unauthorized ||